package com.hrvsr.thread;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class BlockingCountableThreadPool extends ThreadPoolExecutor {
	private static Logger logger = LoggerFactory.getLogger(BlockingCountableThreadPool.class);
	private AtomicInteger threadsAlive = new AtomicInteger(0);
	private int poolSize;
	private Lock lock = new ReentrantLock();
	private Condition canExe = lock.newCondition();
	private int waitInterval = 1000; // milliseconds

	public BlockingCountableThreadPool(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
			BlockingQueue<Runnable> workQueue) {
		super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
	}

	public BlockingCountableThreadPool(int poolSize) {
		super(poolSize, poolSize, 100, TimeUnit.MILLISECONDS, new LinkedBlockingDeque<Runnable>());
		this.poolSize = poolSize;
	}

	@Override
	public int getActiveCount() {
		return threadsAlive.get();
	}

	@Override
	protected void afterExecute(Runnable r, Throwable t) {
		unlock();
		threadsAlive.decrementAndGet();
	}

	private void unlock() {
		lock.lock();
		try {
			canExe.signalAll();
		} finally {
			lock.unlock();
		}
	}

	@Override
	public void execute(final Runnable command) {
		while (threadsAlive.get() >= poolSize) {
			waitForPosition();
		}
		threadsAlive.incrementAndGet();
		super.execute(command);
	}

	private void waitForPosition() {
		lock.lock();
		try {
			canExe.await(waitInterval, TimeUnit.MILLISECONDS);
		} catch (InterruptedException e) {
			logger.error("subscriber thread interrupted while waiting for position", e);
		} finally {
			lock.unlock();
		}
	}

	@Override
	public void shutdown() {
		while(this.getQueue().size()>0){
			waitForPosition();
		}
		super.shutdown();
	}
}
